import os
import cv2
import sys

import numpy as np
from PIL import ImageFont, Image, ImageDraw
from tqdm import tqdm

current_dir = os.path.dirname(__file__)
work_path = os.path.dirname(os.path.dirname(current_dir))

sys.path.append(work_path)
sys.path.append(os.path.join(work_path, "ocr"))

from ocr.document_images_classification import DocumentImagesClassification


def get_paths(src_dirs, ext):
    paths = []
    if os.path.isfile(src_dirs):
        with open(src_dirs, "r") as fr:
            for line in fr.readlines():
                paths.append(line.strip())
        return paths

    for dirs, sub_folder, files in os.walk(src_dirs):
        for name in files:
            if os.path.splitext(name)[-1].lower() in ext:
                paths.append(os.path.join(dirs, name))

    paths.sort(key=lambda x: x.lower())

    return paths


def main():
    image_classifier = DocumentImagesClassification("/Users/liubanggui/Documents/projects/service/document-images-classification")
    image_paths = get_paths("/Users/liubanggui/Documents/datasets/classification/kuaipei/medical-report/health-doc-classify-test/门急诊病历_套打手写", [".jpg", ".jpeg", ".png"])
    # font = ImageFont.truetype("/data/projects/fonts/SimHei.ttf", 70, encoding="utf-8")
    save_dir = "outputs"
    os.makedirs(save_dir, exist_ok=True)

    for path in tqdm(image_paths[:]):
        # print(path)
        ext = os.path.splitext(path)[-1]
        name = os.path.basename(path)

        with open(path, "rb") as f:
            img = f.read()
        img = np.fromstring(img, np.uint8)

        results = image_classifier.process(img)
        print(results.results)

        # img = cv2.imread(path)
        # # print(path, results)
        # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        # w, h = font.getsize("我的")
        # img = Image.fromarray(img)
        # width, height = img.size
        # draw = ImageDraw.Draw(img)
        # try:
        #     for i, res in enumerate(results["results"]["member_information"]):
        #         if "household_name" in res:
        #             text = res["household_name"]
        #         elif "name" in res:
        #             text = f"{res['name']}|{res['relation']}"
        #
        #         draw.text((100, height - (i + 1) * h), text, (0, 255, 0), font)
        # except Exception as e:
        #     print(path)
        #
        # img.save(os.path.join(save_dir, os.path.basename(path)))


if __name__ == '__main__':
    main()
